# 导入所需的模块和类
from langchain_core.prompts import ChatPromptTemplate
from langchain_community.llms.tongyi import Tongyi
from langchain.schema import StrOutputParser
from langchain_core.runnables import RunnablePassthrough

# 创建一个 Tongyi 实例，这是一个语言模型（LLM）的接口
llm = Tongyi()

# 定义四个 ChatPromptTemplate 对象，用于生成不同问题
prompt1 = ChatPromptTemplate.from_template(
    "生成一个{attribute}属性的颜色。除了返回这个颜色的名字不要做其他事:"
)  # 用于生成询问某种属性的颜色的问题
prompt2 = ChatPromptTemplate.from_template(
    "什么水果是这个颜色:{color},只返回这个水果的名字不要做其他事情:"
)  # 用于生成询问某种颜色的水果的问题
prompt3 = ChatPromptTemplate.from_template(
    "哪个国家的国旗有这个颜色:{color},只返回这个国家的名字不要做其他事情:"
)  # 用于生成询问某种颜色的国旗的国家的问题
prompt4 = ChatPromptTemplate.from_template(
    "请问{country}有{fruit}吗？"
)  # 用于生成询问某个国家是否有某种水果的问题

# 定义一个模型解析器，用于解析语言模型的输出
model_parser = llm | StrOutputParser()

# 定义一个链，用于生成一个颜色
chain_color_generator = prompt1 | model_parser

# 定义两个链，用于将颜色转换为水果和国家
chain_color_to_fruit = prompt2 | model_parser
chain_color_to_country = prompt3 | model_parser

# 定义一个链，用于生成一个完整的问题
chain_question_generator = (
    chain_color_generator | {
        "color": RunnablePassthrough(),
        "fruit": chain_color_to_fruit,
        "country": chain_color_to_country
    } | prompt4 | model_parser
)

# 调用 chain_question_generator 的 invoke 方法，并传入 "红色" 来生成问题并获取答案
answer = chain_question_generator.invoke({"attribute": "红色"})

# 打印答案
print(answer)
